#!/usr/bin/env python

""" 
使用 SO100Leader 和 SO100Follower 完成抓取任务

```
from lerobot.robots.so100_leader import SO100Leader, SO100LeaderConfig
from lerobot.robots.so100_follower import SO100Follower, SO100FollowerConfig
from lerobot.cameras import CameraConfig
```

```
# 主臂配置
leader_config = SO100LeaderConfig(port="/dev/ttyUSB0")
leader = SO100Leader(leader_config)
leader.connect()
```

```
# 从臂配置
follower_config = SO100FollowerConfig(
    port="/dev/ttyUSB1",
    cameras={"main": CameraConfig(type="usb", width=640, height=480)}
)
follower = SO100Follower(follower_config)
follower.connect()
```

```
# 远程操作 
while True:
    action = leader.get_action()  # 从主臂读取动作
        follower.send_action(action)  # 控制从臂执行动作
        obs = follower.get_observation()  # 获取从臂观测（关节位置和图像）
        # 可记录 action 和 obs 到数据集
```



```
# 注意事项：
    无相机支持：
    SO100Leader 不包含相机，依赖从臂（如 SO100Follower）提供视觉反馈。
    如果任务需要主臂视觉引导，需修改代码添加相机支持。
```
"""



# 用于作为“主臂”（leader arm）进行远程操作（teleoperation），通过读取主臂的关节位置来生成动作指令，
# 供从臂（follower arm，例如 SO100Follower）或机器人学习系统使用
import logging
import time

from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
from lerobot.motors import Motor, MotorCalibration, MotorNormMode
from lerobot.motors.feetech import (
    FeetechMotorsBus,
    OperatingMode,
)

from ..teleoperator import Teleoperator
from .config_so100_leader import SO100LeaderConfig

logger = logging.getLogger(__name__)


class SO100Leader(Teleoperator):
    """
        读取主臂关节位置：通过 get_action 方法，获取主臂当前各个关节的位置（归一化值），作为动作指令输出。

        无反馈机制：当前版本的 send_feedback 方法未实现（抛出 NotImplementedError），表明不支持向主臂发送力反馈或其他反馈信号。
    """

    config_class = SO100LeaderConfig
    name = "so100_leader"

    def __init__(self, config: SO100LeaderConfig):
        super().__init__(config)
        self.config = config
        self.bus = FeetechMotorsBus(
            port=self.config.port,
            motors={
                "shoulder_pan": Motor(1, "sts3215", MotorNormMode.RANGE_M100_100),
                "shoulder_lift": Motor(2, "sts3215", MotorNormMode.RANGE_M100_100),
                "elbow_flex": Motor(3, "sts3215", MotorNormMode.RANGE_M100_100),
                "wrist_flex": Motor(4, "sts3215", MotorNormMode.RANGE_M100_100),
                "wrist_roll": Motor(5, "sts3215", MotorNormMode.RANGE_M100_100),
                "gripper": Motor(6, "sts3215", MotorNormMode.RANGE_0_100),
            },
            calibration=self.calibration,
        )

    @property
    def action_features(self) -> dict[str, type]:
        return {f"{motor}.pos": float for motor in self.bus.motors}

    @property
    def feedback_features(self) -> dict[str, type]:
        return {}

    @property
    def is_connected(self) -> bool:
        return self.bus.is_connected

    def connect(self, calibrate: bool = True) -> None:
        if self.is_connected:
            raise DeviceAlreadyConnectedError(f"{self} already connected")

        self.bus.connect()
        if not self.is_calibrated and calibrate:
            self.calibrate()

        self.configure()
        logger.info(f"{self} connected.")

    @property
    def is_calibrated(self) -> bool:
        return self.bus.is_calibrated

    def calibrate(self) -> None:
        """
            校准过程：
                1.禁用电机扭矩，设置位置控制模式。
                2.用户将机械臂移动到运动范围中间，按回车记录归位偏移（homing_offsets）。
                3.除 wrist_roll（全旋转关节，范围固定为 0–4095）外，逐一移动其他关节，记录最小和最大位置。
                4.生成 MotorCalibration 对象，保存校准数据到文件（calibration_fpath）。
                5.用途：确保主臂的关节位置与实际硬件运动范围对齐。
        """
        logger.info(f"\nRunning calibration of {self}")
        self.bus.disable_torque()
        for motor in self.bus.motors:
            self.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)

        input(f"Move {self} to the middle of its range of motion and press ENTER....")
        homing_offsets = self.bus.set_half_turn_homings()

        full_turn_motor = "wrist_roll"
        unknown_range_motors = [motor for motor in self.bus.motors if motor != full_turn_motor]
        print(
            f"Move all joints except '{full_turn_motor}' sequentially through their "
            "entire ranges of motion.\nRecording positions. Press ENTER to stop..."
        )
        range_mins, range_maxes = self.bus.record_ranges_of_motion(unknown_range_motors)
        range_mins[full_turn_motor] = 0
        range_maxes[full_turn_motor] = 4095

        self.calibration = {}
        for motor, m in self.bus.motors.items():
            self.calibration[motor] = MotorCalibration(
                id=m.id,
                drive_mode=0,
                homing_offset=homing_offsets[motor],
                range_min=range_mins[motor],
                range_max=range_maxes[motor],
            )

        self.bus.write_calibration(self.calibration)
        self._save_calibration()
        print(f"Calibration saved to {self.calibration_fpath}")

    def configure(self) -> None:
        self.bus.disable_torque()
        self.bus.configure_motors()
        for motor in self.bus.motors:
            self.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)

    def setup_motors(self) -> None:
        for motor in reversed(self.bus.motors):
            input(f"Connect the controller board to the '{motor}' motor only and press enter.")
            self.bus.setup_motor(motor)
            print(f"'{motor}' motor id set to {self.bus.motors[motor].id}")
    
    # 读取关节位置
    # 返回动作字典。动作格式：{"shoulder_pan.pos": float, ..., "gripper.pos": float}。
    def get_action(self) -> dict[str, float]:
        start = time.perf_counter()
        action = self.bus.sync_read("Present_Position")
        action = {f"{motor}.pos": val for motor, val in action.items()}
        dt_ms = (time.perf_counter() - start) * 1e3
        logger.debug(f"{self} read action: {dt_ms:.1f}ms")
        return action

    # 无反馈机制
    def send_feedback(self, feedback: dict[str, float]) -> None:
        # TODO(rcadene, aliberts): Implement force feedback
        raise NotImplementedError

    def disconnect(self) -> None:
        if not self.is_connected:
            DeviceNotConnectedError(f"{self} is not connected.")

        self.bus.disconnect()
        logger.info(f"{self} disconnected.")
